# 线性调频LFM雷达数据集
import os
import random
from typing import Tuple
import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor

class LfmDss(Dataset):
    ds_base = 'D:/research/pykan/work/datasets'
    ds_pt_X = './work/datasets/irhw_dss/train_X.pt'
    ds_pt_y = './work/datasets/irhw_dss/train_y.pt'

    def __init__(self, ds_fn:str = 'irhw_dss/labels.txt'):
        self.ds_fn = f'{LfmDss.ds_base}/{ds_fn}'
        if os.path.exists(LfmDss.ds_pt_X):
            self.X = torch.load(LfmDss.ds_pt_X, weights_only=True)
            self.y = torch.load(LfmDss.ds_pt_y, weights_only=True)
        else:
            self.X, self.y = self.load_ds()
            torch.save(self.X, LfmDss.ds_pt_X)
            torch.save(self.y, LfmDss.ds_pt_y)
        print(f'self.X: {self.X.shape}; self.y: {self.y.shape};')

    def __len__(self):
        return len(self.X)
    
    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]
    
    def load_ds(self) -> Tuple[torch.tensor, torch.tensor]:
        # self.ds_label = f'{self.ds_base}/irhw_dss/labels.txt'
        X = None
        y = None
        with open(self.ds_fn, 'r', encoding='utf-8') as fd:
            for row in fd:
                row = row.strip()
                arrs = row.split(',')
                fn = arrs[0]
                yi = float(arrs[1])
                sample_fn = f'{self.ds_base}/{fn}'
                Xi = self.__load_sample(sample_fn=sample_fn)
                if X is None:
                    X = Xi
                    y = torch.tensor([yi])
                else:
                    X = torch.vstack((X, Xi))
                    y = torch.vstack((y, torch.tensor([yi])))
        return X, y
    
    def __load_sample(self, sample_fn) -> torch.tensor:
        with open(sample_fn, 'r', encoding='utf-8') as fd:
            for row in fd:
                row = row.strip()
                arrs = row.split(',')
                # 因为文件内容只有一行
                x_raw = [float(item) for item in arrs]
                Xi = torch.tensor(x_raw, dtype=torch.float32)
        return Xi
    
    @staticmethod
    def filter_ds(src_ds:str, dst_ds:str, ratio:float = 0.1) -> None:
        with open(src_ds, 'r', encoding='utf-8') as rfd:
            with open(dst_ds, 'w', encoding='utf-8') as wfd:
                for row in rfd:
                    row = row.strip()
                    rv = random.randint(1, 10)
                    if rv == 1: # 仅保留10%
                        wfd.write(f'{row}\n')